# app.py
import time
import os
import json
import subprocess
from typing import Optional, Dict, Any, Deque
from collections import defaultdict, deque

import requests
from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.responses import StreamingResponse

# ===== Ark Runtime（V3）=====
from volcenginesdkarkruntime import Ark

# ===== 导入合同审查模块 =====
from 合同审查 import ContractAnalyzer, classify_contract_type

# ===== 导入工资单分析模块 =====
from 工资单分析 import PayslipAnalyzer, classify_payslip_type
from 考勤记录分析 import AttendanceAnalyzer, classify_attendance_type, extract_attendance_info

# ===== 导入工伤鉴定书分析模块 =====
from 工伤鉴定书分析 import InjuryAssessmentAnalyzer, classify_injury_assessment_type

# ===== 导入劳动赔偿计算器模块 =====
from 劳动赔偿计算器 import LaborCompensationCalculator

# =============== 1) 初始化豆包（Ark）客户端 ===============
ARK_API_KEY ="10a3b378-d38e-4097-b2ce-c4575ff09cb6"
# 如果需要自定义 base_url（企业私有化或网关），可在 Ark(...) 里传 base_url=
client = Ark(api_key=ARK_API_KEY)

# 你的推理接入点（Endpoint ID），到方舟控制台创建，选择带 Function Calling 能力的模型（如 doubao-pro 32k functioncall）
# 注意：不同租户/空间 Endpoint ID 会不一样。
ENDPOINT_ID ="doubao-seed-1-6-250615"

# =============== 2) 初始化合同审查器 ===============
CONTRACT_API_KEY = "sk-20856422ed6644e3827b9d5403c9542a"  # DeepSeek API密钥
contract_analyzer = ContractAnalyzer(CONTRACT_API_KEY)

# =============== 3) 初始化工资单分析器 ===============
PAYROLL_API_KEY = "sk-20856422ed6644e3827b9d5403c9542a"  # DeepSeek API密钥（与合同审查共用）
payroll_analyzer = PayslipAnalyzer(PAYROLL_API_KEY)

# ===== 初始化考勤记录分析器 =====
ATTENDANCE_API_KEY = "sk-20856422ed6644e3827b9d5403c9542a"  # DeepSeek API密钥（共用）
attendance_analyzer = AttendanceAnalyzer(ATTENDANCE_API_KEY)

# ===== 初始化工伤鉴定书分析器 =====
INJURY_ASSESSMENT_API_KEY = "sk-20856422ed6644e3827b9d5403c9542a"  # DeepSeek API密钥（共用）
injury_assessment_analyzer = InjuryAssessmentAnalyzer(INJURY_ASSESSMENT_API_KEY)

# ===== 初始化劳动赔偿计算器 =====
labor_compensation_calculator = LaborCompensationCalculator()

# =============== 会话内存（多轮） ===============
# 每个 session_id -> {"history": deque[[user, assistant], ...], "slots": {"location":..., "date":...}}
MEM: Dict[str, Dict[str, Any]] = defaultdict(lambda: {
    "history": deque(maxlen=20),  # 最近20轮
    "slots": {}
})

# ========= 工具函数 =========
def gradio_history_to_messages(history, system_prompt=None, latest_user=None):
    msgs = []
    if system_prompt:
        msgs.append({"role": "system", "content": system_prompt})
    if history:
        for pair in history:
            if not isinstance(pair, (list, tuple)) or len(pair) != 2:
                continue
            u, a = pair
            if u:
                msgs.append({"role": "user", "content": u})
            if a:
                msgs.append({"role": "assistant", "content": a})
    if latest_user:
        msgs.append({"role": "user", "content": latest_user})
    return msgs

def backfill_slots(intent: str, route: Dict[str, Any], session_slots: Dict[str, Any]) -> Dict[str, Any]:
    # 仅对需要槽位的意图做回填
    if intent == "weather":
        if not route.get("location"):
            route["location"] = session_slots.get("location")
        if not route.get("date"):
            route["date"] = session_slots.get("date")
    if intent == "labor_legal":
        # 示例：可扩展 labor_legal 的法域回填
        if not route.get("jurisdiction"):
            route["jurisdiction"] = session_slots.get("jurisdiction")
    if intent == "contract_review":
        # 合同审查：回填文件路径
        if not route.get("file_path"):
            route["file_path"] = session_slots.get("file_path")
    if intent == "labor_compensation":
        # 劳动赔偿计算：回填参数
        if not route.get("compensation_type"):
            route["compensation_type"] = session_slots.get("compensation_type")
        if not route.get("monthly_salary"):
            route["monthly_salary"] = session_slots.get("monthly_salary")
        if not route.get("work_years"):
            route["work_years"] = session_slots.get("work_years")
        if not route.get("city"):
            route["city"] = session_slots.get("city")
    return route

def update_slots_after_success(intent: str, route: Dict[str, Any], session_slots: Dict[str, Any]):
    if intent == "weather":
        if route.get("location"):
            session_slots["location"] = route["location"]
        if route.get("date"):
            session_slots["date"] = route["date"]
    if intent == "labor_legal":
        if route.get("jurisdiction"):
            session_slots["jurisdiction"] = route["jurisdiction"]
    if intent == "contract_review":
        if route.get("file_path"):
            session_slots["file_path"] = route["file_path"]
    if intent == "labor_compensation":
        if route.get("compensation_type"):
            session_slots["compensation_type"] = route["compensation_type"]
        if route.get("monthly_salary"):
            session_slots["monthly_salary"] = route["monthly_salary"]
        if route.get("work_years"):
            session_slots["work_years"] = route["work_years"]
        if route.get("city"):
            session_slots["city"] = route["city"]

# =============== 2) 路由工具Schema ===============
ROUTER_TOOL = {
    "type": "function",
    "function": {
        "name": "route_user_query",
        "description": "对用户问题进行意图识别与槽位抽取，只返回应该调用的工具类型和关键字段。对于劳动赔偿计算，需要从用户问题中提取：赔偿类型（经济补偿金、赔偿金、双倍工资等）、月工资数额、工作年限、所在城市。",
        "parameters": {
            "type": "object",
            "properties": {
                "intent": {
                    "type": "string",
                    "enum": ["weather", "labor_legal", "contract_review", "payroll_analysis", "attendance_analysis", "labor_compensation", "general_qa", "unknown"],
                    "description": "根据用户问题选择的工具：weather=查询天气, labor_legal=劳动维权咨询, contract_review=合同审查, payroll_analysis=工资单分析, attendance_analysis=考勤记录分析, labor_compensation=劳动赔偿计算, general_qa=常识咨询, unknown=无法判断"
                },
                "location": {"type": "string", "description": "与天气或场景相关的地点"},
                "date": {"type": "string", "description": "与天气相关的日期，建议YYYY-MM-DD"},
                "question": {"type": "string", "description": "抽取后的简要问题"},
                "jurisdiction": {"type": "string", "description": "劳动维权适用法域/地区"},
                "file_path": {"type": "string", "description": "合同文件路径（支持PDF、DOCX或图片格式）"},
                "payroll_file_path": {"type": "string", "description": "工资单文件路径（支持PDF、DOCX或图片格式）"},
                "attendance_file_path": {"type": "string", "description": "考勤记录文件路径（支持PDF、DOCX或图片格式）"},
                "compensation_type": {"type": "string", "description": "赔偿类型：经济补偿金、赔偿金、双倍工资、加班费等"},
                "monthly_salary": {"type": "number", "description": "月工资（元）"},
                "work_years": {"type": "number", "description": "工作年限（年）"},
                "city": {"type": "string", "description": "所在城市"}
            },
            "required": ["intent"]
        }
    }
}

# =============== 3) 业务工具 ===============
def query_weather(location: Optional[str], date: Optional[str]) -> str:
    if not location:
        return "我需要一个地点才能查询天气：请提供城市或地名。"
    city_db = {
        "北京": (39.9042, 116.4074),
        "上海": (31.2304, 121.4737),
        "广州": (23.1291, 113.2644),
        "深圳": (22.5431, 114.0579),
        "杭州": (30.2741, 120.1551),
    }
    lat, lon = city_db.get(location, (31.2304, 121.4737))  # 默认上海
    url = "https://api.open-meteo.com/v1/forecast"
    params = {
        "latitude": lat,
        "longitude": lon,
        "hourly": "temperature_2m,precipitation",
        "forecast_days": 1
    }
    try:
        resp = requests.get(url, params=params, timeout=8)
        resp.raise_for_status()
        data = resp.json()
        temps = data.get("hourly", {}).get("temperature_2m", [])
        if not temps:
            return f"{location} 的天气数据暂不可用。"
        avg_temp = sum(temps) / len(temps)
        suffix = f"（你也可以直接问：明天呢/后天呢？我会记住地点{location}）"
        return f"{location} 近况（示例数据）：平均气温约 {avg_temp:.1f}°C。{suffix}"
    except Exception as e:
        return f"查询天气失败（示例接口）：{e}"

def labor_legal_consult(question: str) -> str:
    messages = [
        {"role": "system", "content": "你是恒纪元劳动法小助手，用电话沟通的口语化风格回复，语气亲切自然，避免书面化表达。"},
        {"role": "user", "content": question}
    ]
    try:
        completion = client.chat.completions.create(
            model=ENDPOINT_ID,
            messages=messages,
            tool_choice="none",
            thinking={"type": "disabled"}
        )
        if not completion.choices:
            return "模型没有返回内容，请稍后再试。"
        return completion.choices[0].message.content
    except Exception as e:
        return f"常识问答调用失败：{e}"

def general_qa(question: str) -> str:
    messages = [
        {"role": "system", "content": "你是一个靠谱、简洁的中文助理。用电话沟通的口语化风格回复，语气亲切自然，避免书面化表达。"},
        {"role": "user", "content": question}
    ]
    try:
        completion = client.chat.completions.create(
            model=ENDPOINT_ID,
            messages=messages,
            tool_choice="none",
            thinking={"type": "disabled"}
        )
        if not completion.choices:
            return "模型没有返回内容，请稍后再试。"
        return completion.choices[0].message.content
    except Exception as e:
        return f"常识问答调用失败：{e}"

def contract_review_consult(file_path: str) -> str:
    """合同审查工具函数"""
    if not file_path:
        return "请提供合同文件路径进行审查。支持PDF、DOCX或图片格式。"
    
    if not os.path.exists(file_path):
        return f"文件不存在：{file_path}。请检查文件路径是否正确。"
    
    try:
        # 提取文本
        text = contract_analyzer.extract_text(file_path)
        if not text.strip():
            return "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
        
        # 识别合同类型
        contract_type = classify_contract_type(text)
        
        # 进行合同分析（非流式，返回完整结果）
        result_parts = []
        for chunk in contract_analyzer.analyze_contract_stream(text, contract_type, [file_path]):
            result_parts.append(chunk)
        
        analysis_result = "".join(result_parts)
        return f"📋 合同类型：{contract_type}\n\n📝 分析结果：\n{analysis_result}"
        
    except Exception as e:
        return f"合同审查失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"

def payroll_analysis_consult(file_path: str) -> str:
    """工资单分析工具函数"""
    if not file_path:
        return "请提供工资单文件路径进行分析。支持PDF、DOCX或图片格式。"
    
    if not os.path.exists(file_path):
        return f"文件不存在：{file_path}。请检查文件路径是否正确。"
    
    try:
        # 提取文本
        text = payroll_analyzer.extract_text(file_path)
        if not text.strip():
            return "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
        
        # 识别工资单类型
        payroll_type = classify_payslip_type(text)
        
        # 提取工资单信息
        from 工资单分析 import extract_payslip_info
        extracted_info = extract_payslip_info(text)
        
        # 进行工资单分析（非流式，返回完整结果）
        result_parts = []
        for chunk in payroll_analyzer.analyze_payslip_stream(text, payroll_type, extracted_info):
            result_parts.append(chunk)
        
        analysis_result = "".join(result_parts)
        return f"💰 工资单类型：{payroll_type}\n\n📊 分析结果：\n{analysis_result}"
        
    except Exception as e:
        return f"工资单分析失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"

def attendance_analysis_consult(file_path: str) -> str:
    """考勤记录分析工具函数"""
    if not file_path:
        return "请提供考勤记录文件路径进行分析。支持PDF、DOCX或图片格式。"
    
    if not os.path.exists(file_path):
        return f"文件不存在：{file_path}。请检查文件路径是否正确。"
    
    try:
        # 提取文本
        text = attendance_analyzer.extract_text(file_path)
        if not text.strip():
            return "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
        
        # 识别考勤记录类型
        attendance_type = classify_attendance_type(text)
        
        # 提取考勤记录信息
        extracted_info = extract_attendance_info(text)
        
        # 进行考勤记录分析（非流式，返回完整结果）
        result_parts = []
        for chunk in attendance_analyzer.analyze_attendance_stream(text, attendance_type, extracted_info):
            result_parts.append(chunk)
        
        analysis_result = "".join(result_parts)
        return f"📋 考勤记录类型：{attendance_type}\n\n📊 分析结果：\n{analysis_result}"
        
    except Exception as e:
        return f"考勤记录分析失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"

def injury_assessment_analysis_consult(file_path: str) -> str:
    """工伤鉴定书分析工具函数"""
    if not file_path:
        return "请提供工伤鉴定书文件路径进行分析。支持PDF、DOCX或图片格式。"
    
    if not os.path.exists(file_path):
        return f"文件不存在：{file_path}。请检查文件路径是否正确。"
    
    try:
        # 提取文本
        text = injury_assessment_analyzer.extract_text(file_path)
        if not text.strip():
            return "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
        
        # 识别工伤鉴定书类型
        injury_type = classify_injury_assessment_type(text)
        
        # 进行工伤鉴定书分析（非流式，返回完整结果）
        result_parts = []
        for chunk in injury_assessment_analyzer.analyze_injury_assessment_stream(text, injury_type, [file_path]):
            result_parts.append(chunk)
        
        analysis_result = "".join(result_parts)
        return f"🏥 工伤鉴定书类型：{injury_type}\n\n📋 分析结果：\n{analysis_result}"
        
    except Exception as e:
        return f"工伤鉴定书分析失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"

def labor_compensation_calculate(compensation_type: str, monthly_salary: float, work_years: float, city: str) -> str:
    """劳动赔偿计算工具函数"""
    if not compensation_type:
        return "请指定赔偿类型（经济补偿金、赔偿金、双倍工资、加班费等）。"
    
    # 检查缺失的参数并提供友好的引导
    missing_params = []
    if not monthly_salary or monthly_salary <= 0:
        missing_params.append("月工资金额")
    if not work_years or work_years <= 0:
        missing_params.append("工作年限")
    if not city:
        missing_params.append("所在城市")
    
    if missing_params:
        return f"💡 要计算{compensation_type}，我还需要以下信息：\n\n" + \
               "\n".join([f"• {param}" for param in missing_params]) + \
               "\n\n请提供这些信息，例如：\n" + \
               "我在北京工作了3年，月工资8000元，想计算经济补偿金。"
    
    try:
        # 获取城市列表，检查城市是否支持
        city_list = labor_compensation_calculator.get_city_list()
        
        # 智能匹配城市名称
        matched_city = None
        if city in city_list:
            matched_city = city
        else:
            # 尝试添加"市"后缀匹配
            city_with_suffix = city + "市" if not city.endswith("市") else city
            if city_with_suffix in city_list:
                matched_city = city_with_suffix
            else:
                # 尝试模糊匹配
                for supported_city in city_list:
                    if city in supported_city or supported_city.replace("市", "") == city:
                        matched_city = supported_city
                        break
        
        if not matched_city:
            return f"暂不支持城市：{city}。支持的城市有：{', '.join(city_list)}"
        
        # 使用匹配到的城市名称
        city = matched_city
        
        # 根据赔偿类型进行计算
        if "经济补偿" in compensation_type or compensation_type == "经济补偿金":
            result = labor_compensation_calculator.calculate_economic_compensation(
                monthly_salary=monthly_salary,
                work_years=work_years,
                city=city
            )
        elif "赔偿金" in compensation_type or "2N" in compensation_type:
            result = labor_compensation_calculator.calculate_severance_pay(
                monthly_salary=monthly_salary,
                work_years=work_years,
                city=city
            )
        elif "双倍工资" in compensation_type:
            result = labor_compensation_calculator.calculate_double_salary(
                monthly_salary=monthly_salary,
                unsigned_months=min(int(work_years * 12), 11)  # 双倍工资最多11个月
            )
        else:
            return f"暂不支持的赔偿类型：{compensation_type}。支持的类型有：经济补偿金、赔偿金、双倍工资。"
        
        # 生成详细报告
        report = labor_compensation_calculator.generate_calculation_report([result])
        
        return report
        
    except Exception as e:
        return f"劳动赔偿计算失败：{str(e)}。请检查输入参数是否正确。"

# =============== 4) 第一段：路由（仅暴露路由工具） ===============
def first_stage_route(user_query: str) -> Dict[str, Any]:
    messages = [
        {"role": "system", "content": "请仅调用 route_user_query 工具，判断该走哪个业务工具并抽取必要字段。"},
        {"role": "user", "content": user_query}
    ]
    resp = client.chat.completions.create(
        model=ENDPOINT_ID,
        messages=messages,
        tools=[ROUTER_TOOL],
        tool_choice="auto",
        thinking={"type": "disabled"}
    )
    choice = resp.choices[0]
    tool_calls = getattr(choice.message, "tool_calls", None)
    if not tool_calls:
        return {"intent": "unknown", "question": user_query}
    call = tool_calls[0]
    if call.function.name != "route_user_query":
        return {"intent": "unknown", "question": user_query}
    args = json.loads(call.function.arguments or "{}")
    return {
        "intent": args.get("intent", "unknown"),
        "location": args.get("location"),
        "date": args.get("date"),
        "question": args.get("question") or user_query,
        "jurisdiction": args.get("jurisdiction"),
        "file_path": args.get("file_path"),
        "payroll_file_path": args.get("payroll_file_path"),
        "attendance_file_path": args.get("attendance_file_path"),
        "compensation_type": args.get("compensation_type"),
        "monthly_salary": args.get("monthly_salary"),
        "work_years": args.get("work_years"),
        "city": args.get("city")
    }

# =============== 5) 第二段：执行选中工具 ===============
def second_stage_execute(router_result: Dict[str, Any]) -> str:
    intent = router_result.get("intent", "unknown")
    print(f"[Router] intent={intent} | route={router_result}")

    if intent == "weather":
        return query_weather(router_result.get("location"), router_result.get("date"))
    if intent == "labor_legal":
        return labor_legal_consult(router_result.get("question", ""))
    if intent == "contract_review":
        return contract_review_consult(router_result.get("file_path", ""))
    if intent == "payroll_analysis":
        return payroll_analysis_consult(router_result.get("payroll_file_path", ""))
    if intent == "attendance_analysis":
        return attendance_analysis_consult(router_result.get("attendance_file_path", ""))
    if intent == "injury_assessment_analysis":
        return injury_assessment_analysis_consult(router_result.get("injury_assessment_file_path", ""))
    if intent == "labor_compensation":
        return labor_compensation_calculate(
            router_result.get("compensation_type", ""),
            router_result.get("monthly_salary", 0),
            router_result.get("work_years", 0),
            router_result.get("city", "")
        )
    if intent == "general_qa":
        return general_qa(router_result.get("question", ""))
    return "我没能判断要用哪个工具。你可以换种问法：例如\"上海明天的天气怎么样？\"、\"加班不给加班费怎么办？\"、\"请审查这份合同：文件路径\"、\"请分析这份工资单：文件路径\"、\"请分析这份考勤记录：文件路径\"、\"请分析这份工伤鉴定书：文件路径\"或\"我在北京工作3年，月薪8000，计算经济补偿金\""

# =============== 6) HTTP 接口 ===============
class ChatIn(BaseModel):
    query: str
    history: Optional[list] = None   # 前端可选传（兼容你现有写法）
    session_id: Optional[str] = "default"  # 新增：多会话支持
    attendance_file_path: Optional[str] = None  # 考勤文件路径
    file_path: Optional[str] = None  # 合同文件路径
    payroll_file_path: Optional[str] = None  # 工资单文件路径
    injury_assessment_file_path: Optional[str] = None  # 工伤鉴定书文件路径

app = FastAPI()

@app.post("/chat_stream")
def chat_stream(inp: ChatIn):
    """
    流式接口：根据工具路由决定如何输出；并把本轮问答写入会话内存。
    """
    def token_generator():
        try:
            sid = inp.session_id or "default"
            mem = MEM[sid]

            # 用最新一轮用户输入做路由；随后做槽位回填
            t0 = time.perf_counter()
            route = first_stage_route(inp.query)
            route = backfill_slots(route.get("intent", "unknown"), route, mem["slots"])
            
            # 从请求中获取文件路径并设置到route中
            if inp.attendance_file_path and not route.get("attendance_file_path"):
                route["attendance_file_path"] = inp.attendance_file_path
            if inp.file_path and not route.get("file_path"):
                route["file_path"] = inp.file_path
            if inp.payroll_file_path and not route.get("payroll_file_path"):
                route["payroll_file_path"] = inp.payroll_file_path
            if inp.injury_assessment_file_path and not route.get("injury_assessment_file_path"):
                route["injury_assessment_file_path"] = inp.injury_assessment_file_path
            
            t1 = time.perf_counter()
            print(f"[PROFILE] tool routing took {(t1 - t0):.4f} s")
            intent = route.get("intent", "unknown")
            print(f"[Router] intent={intent} | route={route}")

            # ========== weather：一次性返回 + 内存追加 ==========
            if intent == "weather":
                answer = query_weather(route.get("location"), route.get("date"))
                # 记录历史 & 槽位
                mem["history"].append([inp.query, answer])
                update_slots_after_success(intent, route, mem["slots"])
                yield f"{answer}\n\n[调试] 选用工具: {intent}"
                return

            # ========== contract_review：流式 ==========
            if intent == "contract_review":
                file_path = route.get("file_path", "")
                if not file_path:
                    error_msg = "请提供合同文件路径进行审查。支持PDF、DOCX或图片格式。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return
                
                if not os.path.exists(file_path):
                    error_msg = f"文件不存在：{file_path}。请检查文件路径是否正确。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return
                
                try:
                    # 提取文本
                    text = contract_analyzer.extract_text(file_path)
                    if not text.strip():
                        error_msg = "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
                        mem["history"].append([inp.query, error_msg])
                        yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                        return
                    
                    # 识别合同类型
                    contract_type = classify_contract_type(text)
                    yield f"📋 合同类型：{contract_type}\n\n📝 分析结果：\n"
                    
                    # 流式分析
                    chunks = []
                    for chunk in contract_analyzer.analyze_contract_stream(text, contract_type, [file_path]):
                        chunks.append(chunk)
                        yield chunk
                    
                    # 记录历史 & 槽位
                    full_answer = f"📋 合同类型：{contract_type}\n\n📝 分析结果：\n{''.join(chunks)}"
                    mem["history"].append([inp.query, full_answer])
                    update_slots_after_success(intent, route, mem["slots"])
                    
                    yield f"\n\n[调试] 选用工具: {intent}"
                    return
                    
                except Exception as e:
                    error_msg = f"合同审查失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return

            # ========== payroll_analysis：流式 ==========
            if intent == "payroll_analysis":
                file_path = route.get("payroll_file_path", "")
                if not file_path:
                    error_msg = "请提供工资单文件路径进行分析。支持PDF、DOCX或图片格式。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return
                
                if not os.path.exists(file_path):
                    error_msg = f"文件不存在：{file_path}。请检查文件路径是否正确。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return
                
                try:
                    # 提取文本
                    text = payroll_analyzer.extract_text(file_path)
                    if not text.strip():
                        error_msg = "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
                        mem["history"].append([inp.query, error_msg])
                        yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                        return
                    
                    # 识别工资单类型
                    payroll_type = classify_payslip_type(text)
                    yield f"💰 工资单类型：{payroll_type}\n\n📊 分析结果：\n"
                    
                    # 提取工资单信息
                    from 工资单分析 import extract_payslip_info
                    extracted_info = extract_payslip_info(text)
                    
                    # 流式分析
                    chunks = []
                    for chunk in payroll_analyzer.analyze_payslip_stream(text, payroll_type, extracted_info):
                        chunks.append(chunk)
                        yield chunk
                    
                    # 记录历史 & 槽位
                    full_answer = f"💰 工资单类型：{payroll_type}\n\n📊 分析结果：\n{''.join(chunks)}"
                    mem["history"].append([inp.query, full_answer])
                    update_slots_after_success(intent, route, mem["slots"])
                    
                    yield f"\n\n[调试] 选用工具: {intent}"
                    return
                    
                except Exception as e:
                    error_msg = f"工资单分析失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return

            # ========== attendance_analysis：流式 ==========
            if intent == "attendance_analysis":
                file_path = route.get("attendance_file_path", "")
                if not file_path:
                    error_msg = "请提供考勤记录文件路径进行分析。支持PDF、DOCX或图片格式。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return
                
                if not os.path.exists(file_path):
                    error_msg = f"文件不存在：{file_path}。请检查文件路径是否正确。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return
                
                try:
                    # 提取文本
                    text = attendance_analyzer.extract_text(file_path)
                    if not text.strip():
                        error_msg = "未能从文件中提取到有效文本内容，请检查文件格式或内容。"
                        mem["history"].append([inp.query, error_msg])
                        yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                        return
                    
                    # 识别考勤记录类型
                    attendance_type = classify_attendance_type(text)
                    yield f"📋 考勤记录类型：{attendance_type}\n\n📊 分析结果：\n"
                    
                    # 提取考勤记录信息
                    extracted_info = extract_attendance_info(text)
                    
                    # 流式分析
                    chunks = []
                    for chunk in attendance_analyzer.analyze_attendance_stream(text, attendance_type, extracted_info):
                        chunks.append(chunk)
                        yield chunk
                    
                    # 记录历史 & 槽位
                    full_answer = f"📋 考勤记录类型：{attendance_type}\n\n📊 分析结果：\n{''.join(chunks)}"
                    mem["history"].append([inp.query, full_answer])
                    update_slots_after_success(intent, route, mem["slots"])
                    
                    yield f"\n\n[调试] 选用工具: {intent}"
                    return
                    
                except Exception as e:
                    error_msg = f"考勤记录分析失败：{str(e)}。请检查文件格式是否支持或联系技术支持。"
                    mem["history"].append([inp.query, error_msg])
                    yield f"{error_msg}\n\n[调试] 选用工具: {intent}"
                    return

            # ========== labor_compensation：返回计算器标识 ==========
            if intent == "labor_compensation":
                # 返回劳动赔偿计算器标识，让前端显示计算表单
                answer = "我来帮您计算劳动赔偿。请在右侧的计算器中填写相关信息："
                
                # 记录历史
                mem["history"].append([inp.query, answer])
                update_slots_after_success(intent, route, mem["slots"])
                
                # 返回带有计算器标识的响应
                import json
                response_data = json.dumps({
                    'content': answer,
                    'type': 'message',
                    'show_calculator': True
                }, ensure_ascii=False)
                yield f"data: {response_data}\n\n"
                return

            # ========== labor_legal / general_qa：流式 ==========
            if intent in ("labor_legal", "general_qa"):
                system_prompt = "用电话沟通的口语化风格回复，语气亲切自然，避免书面化表达。"
                # 以服务端内存为准，如果前端也传history则拼起来（去重很复杂，这里简单合并）
                merged_history = list(mem["history"])
                if inp.history:
                    merged_history.extend(inp.history)

                messages = gradio_history_to_messages(
                    history=merged_history,
                    system_prompt=system_prompt,
                    latest_user=route.get("question") or inp.query
                )

                completion = client.chat.completions.create(
                    model=ENDPOINT_ID,
                    messages=messages,
                    stream=True,
                    tool_choice="none",
                    thinking={"type": "disabled"}
                )

                chunks = []
                with completion:
                    for chunk in completion:
                        delta = getattr(chunk.choices[0].delta, "content", None)
                        if delta:
                            chunks.append(delta)
                            yield delta

                # 结尾：写入历史、更新槽位
                full_answer = "".join(chunks) if chunks else ""
                mem["history"].append([inp.query, full_answer])
                update_slots_after_success(intent, route, mem["slots"])

                yield f"\n\n[调试] 选用工具: {intent}"
                return

            # 兜底
            fallback = "我没能判断要用哪个工具。你可以换种问法：例如“上海明天的天气怎么样？”或“加班不给加班费怎么办？”"
            mem["history"].append([inp.query, fallback])
            yield fallback
            yield f"\n\n[调试] 选用工具: {intent}"

        except Exception as e:
            import traceback, sys
            err = "".join(traceback.format_exception(*sys.exc_info()))
            print("[ERROR] /chat_stream failed:\n", err)
            yield f"后端异常：{e}"

    return StreamingResponse(token_generator(), media_type="text/plain; charset=utf-8")

@app.post("/chat")
def chat_endpoint(inp: ChatIn):
    try:
        sid = inp.session_id or "default"
        mem = MEM[sid]

        route = first_stage_route(inp.query)
        route = backfill_slots(route.get("intent", "unknown"), route, mem["slots"])

        answer = second_stage_execute(route)

        # 记录历史与槽位
        mem["history"].append([inp.query, answer])
        update_slots_after_success(route.get("intent", "unknown"), route, mem["slots"])

        return {
            "router": route,
            "answer": f"{answer}\n\n[调试] 选用工具: {route.get('intent')}",
            "memory": {
                "turns": len(mem["history"]),
                "slots": mem["slots"]
            }
        }
    except Exception as e:
        import traceback, sys
        err = "".join(traceback.format_exception(*sys.exc_info()))
        print("[ERROR] /chat failed:\n", err)
        return {"router": {"intent": "error"}, "answer": f"后端异常：{e}\n{err}"}

@app.post("/reset_session")
def reset_session(session_id: Optional[str] = "default"):
    MEM.pop(session_id or "default", None)
    return {"ok": True}

if __name__ == "__main__":
    # 运行： uvicorn app:app --reload --port 8001
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8001)
